package me.zhengjie.deepseek.service;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import me.zhengjie.deepseek.dto.*;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.util.retry.Retry;

import java.time.Duration;
import java.util.List;
import java.util.Map;

@Service
public class DeepSeekClientService {
    private final WebClient webClient;

    public DeepSeekClientService(WebClient.Builder builder,
                                 @Value("${deepseek.api.base-url}") String baseUrl,
                                 @Value("${deepseek.api.key}") String apiKey) {
        this.webClient = builder
                .baseUrl(baseUrl)
                .defaultHeader(HttpHeaders.AUTHORIZATION, "Bearer " + apiKey)
                .defaultHeader(HttpHeaders.ACCEPT, MediaType.TEXT_EVENT_STREAM_VALUE)
                .build();
    }

    public Flux<String> generateText(List<Message> messages) {
        Map<String, Object> requestBody = Map.of(
                "model", "deepseek-chat",
                "stream", true,
                "messages", messages
        );

        return webClient.post()
                .uri("/chat/completions")
                .bodyValue(requestBody)
                .retrieve()
                .bodyToFlux(String.class)
                .map(chunk -> {
                    // 移除 SSE 前缀 "data: "
                    if (chunk.startsWith("data: ")) {
                        chunk = chunk.substring(6);
                    }

                    // 过滤掉可能的心跳消息
                    if (chunk.equals("[DONE]")) {
                        return "";
                    }

                    try {
                        // 解析 JSON 响应
                        ObjectMapper mapper = new ObjectMapper();
                        JsonNode jsonNode = mapper.readTree(chunk);

                        // 提取内容部分
                        if (jsonNode.has("choices") && jsonNode.get("choices").isArray() &&
                                !jsonNode.get("choices").isEmpty()) {

                            JsonNode choice = jsonNode.get("choices").get(0);
                            // 根据模型返回的结构提取内容
                            if (choice.has("delta") && choice.get("delta").has("content")) {
                                // 对于流式响应，内容在 delta.content
                                return choice.get("delta").get("content").asText();
                            } else if (choice.has("message") && choice.get("message").has("content")) {
                                // 对于完整响应，内容在 message.content
                                return choice.get("message").get("content").asText();
                            }
                        }
                        return ""; // 如果没有找到内容，返回空字符串
                    } catch (Exception e) {
                        System.err.println("Error parsing JSON: " + e.getMessage());
                        return ""; // 解析错误时返回空字符串
                    }
                })
                .filter(content -> !content.isEmpty())
                .concatWithValues("[DONE]") // 发送结束信号
                .timeout(Duration.ofSeconds(60));
    }

    public Mono<AnalyzeResponse> analyzeText(AnalyzeRequest req) {
        return webClient.post()
                .uri("/text/analysis")
                .bodyValue(req)
                .retrieve()
                .bodyToMono(AnalyzeResponse.class)
                .timeout(Duration.ofSeconds(10));
    }
}
